import torch


class AnchorGenerator(object):

    def __init__(self, base_size, scales, ratios, scale_major=True, ctr=None, anchor_strides_fraction=None):
        # in fact , base size is just the anchor stride
        self.base_size = base_size
        self.scales = torch.Tensor(scales)
        self.ratios = torch.Tensor(ratios)
        self.anchor_strides_fraction = anchor_strides_fraction
        if self.anchor_strides_fraction is not None:
            assert len(self.anchor_strides_fraction) == len(scales)

        self.scale_major = scale_major
        self.ctr = ctr
        self.base_anchors = self.gen_base_anchors()

    @property
    def num_base_anchors(self):
        return self.base_anchors.size(0)

    def gen_base_anchors(self):
        w = self.base_size
        h = self.base_size
        if self.ctr is None:
            x_ctr = 0.5 * (w - 1)
            y_ctr = 0.5 * (h - 1)
        else:
            x_ctr, y_ctr = self.ctr

        h_ratios = torch.sqrt(self.ratios)
        w_ratios = 1 / h_ratios

        if self.anchor_strides_fraction is not None:
            if self.scale_major:
                ws = (w * w_ratios[:, None] * self.scales[None, :])
                hs = (h * h_ratios[:, None] * self.scales[None, :])
            else:
                ws = (w * self.scales[:, None] * w_ratios[None, :])
                hs = (h * self.scales[:, None] * h_ratios[None, :])
            base_anchors_list = []

            for j in range(len(self.anchor_strides_fraction)):
                if self.scale_major:
                    per_scale_ws = ws[:, j]
                    per_scale_hs = hs[:, j]
                else:
                    per_scale_ws = ws[j, :]
                    per_scale_hs = hs[j, :]
                total_stride = self.base_size
                partial_stride = total_stride // self.anchor_strides_fraction[j]
                partial_stride_list = [k * partial_stride for k in
                                       range(-self.anchor_strides_fraction[j], self.anchor_strides_fraction[j])]
                partial_stride_list = [k for k in partial_stride_list if
                                       k <= total_stride // 2 and k > -total_stride // 2]
                for x_i in partial_stride_list:
                    for y_i in partial_stride_list:
                        base_anchors = torch.stack(
                            [
                                x_ctr + x_i - 0.5 * (per_scale_ws - 1), y_ctr + y_i - 0.5 * (per_scale_hs - 1),
                                x_ctr + x_i + 0.5 * (per_scale_ws - 1), y_ctr + y_i + 0.5 * (per_scale_hs - 1)
                            ],
                            dim=-1).round()
                        base_anchors_list.append(base_anchors)
            base_anchors = torch.cat(base_anchors_list, dim=0)
        else:
            if self.scale_major:
                ws = (w * w_ratios[:, None] * self.scales[None, :]).view(-1)
                hs = (h * h_ratios[:, None] * self.scales[None, :]).view(-1)
            else:
                ws = (w * self.scales[:, None] * w_ratios[None, :]).view(-1)
                hs = (h * self.scales[:, None] * h_ratios[None, :]).view(-1)

            base_anchors = torch.stack(
                [
                    x_ctr - 0.5 * (ws - 1), y_ctr - 0.5 * (hs - 1),
                    x_ctr + 0.5 * (ws - 1), y_ctr + 0.5 * (hs - 1)
                ],
                dim=-1).round()
        return base_anchors

    def _meshgrid(self, x, y, row_major=True):
        xx = x.repeat(len(y))
        yy = y.view(-1, 1).repeat(1, len(x)).view(-1)
        if row_major:
            return xx, yy
        else:
            return yy, xx

    def grid_anchors(self, featmap_size, stride=16, device='cuda'):
        base_anchors = self.base_anchors.to(device)

        feat_h, feat_w = featmap_size
        shift_x = torch.arange(0, feat_w, device=device) * stride
        shift_y = torch.arange(0, feat_h, device=device) * stride
        shift_xx, shift_yy = self._meshgrid(shift_x, shift_y)
        shifts = torch.stack([shift_xx, shift_yy, shift_xx, shift_yy], dim=-1)
        shifts = shifts.type_as(base_anchors)
        # first feat_w elements correspond to the first row of shifts
        # add A anchors (1, A, 4) to K shifts (K, 1, 4) to get
        # shifted anchors (K, A, 4), reshape to (K*A, 4)

        all_anchors = base_anchors[None, :, :] + shifts[:, None, :]
        all_anchors = all_anchors.view(-1, 4)
        # first A rows correspond to A anchors of (0, 0) in feature map,
        # then (0, 1), (0, 2), ...
        return all_anchors

    def valid_flags(self, featmap_size, valid_size, device='cuda'):
        feat_h, feat_w = featmap_size
        valid_h, valid_w = valid_size
        assert valid_h <= feat_h and valid_w <= feat_w
        valid_x = torch.zeros(feat_w, dtype=torch.uint8, device=device)
        valid_y = torch.zeros(feat_h, dtype=torch.uint8, device=device)
        valid_x[:valid_w] = 1
        valid_y[:valid_h] = 1
        valid_xx, valid_yy = self._meshgrid(valid_x, valid_y)
        valid = valid_xx & valid_yy
        valid = valid[:, None].expand(
            valid.size(0), self.num_base_anchors).contiguous().view(-1)
        return valid
